package day005;

public class Code01_Trie {

    public static class Trie{
        public int pass;
        public int end;
        public Trie[] nexts;

        public Trie(){
            pass = 0;
            end = 0;
            nexts = new Trie[26];
        }

        public void insert(String word){
            if (word == null){
                return;
            }
            char[] chs = word.toCharArray();
            Trie node = this;
            node.pass++;
            int index = 0;
            for(int i = 0; i < chs.length; i++){
                index = chs[i] - 'a';
                if(node.nexts[index] == null){
                    node.nexts[index] = new Trie();
                }
                node = node.nexts[index];
                node.pass++;
            }
            node.end++;
        }

        // word 这个单词加入过几次
        public int search(String word){
            if(word == null){
                return 0;
            }
            char[] chs = word.toCharArray();
            Trie node = this;
            int index = 0;
            for(int i = 0; i <chs.length; i++){
                index = chs[i] - 'a';
                if(node.nexts[index] == null){
                    return 0;
                }
                node = node.nexts[index];
            }
            return node.end;
        }

        // 所有字符串中，以pre为前缀的数量
        public int searchPrefix(String pre){
            if(pre == null){
                return 0;
            }
            char[] chs = pre.toCharArray();
            Trie node = this;
            int index = 0;
            for(int i = 0; i < chs.length; i++){
                index = chs[i] - 'a'; // index 就是只当前遍历到的节点
                if(node.nexts[index] == null){
                    return 0;
                }
                node = node.nexts[index];
            }
            return node.pass;
        }
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("abc");
        trie.insert("bck");
        trie.insert("abd");
        trie.insert("ace");
        trie.insert("abc");
        System.out.println(trie.search("bck"));
        System.out.println(trie.searchPrefix("a"));
    }
}
